package com.jiazhong.spring.ai.elasticsearch.service.impl;

import com.jiazhong.spring.ai.elasticsearch.service.OllamaService;
import jakarta.annotation.Resource;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.stream.Collectors;

@Service
public class OllamaServiceImpl implements OllamaService {
    @Resource
    private ChatClient ollamaChatClient;
    @Resource
    private VectorStore vectorStore;

    @Override
    public Flux<String> stream(String question) {
        // 1. 检索向量数据库
        SearchRequest searchRequest = SearchRequest.builder()
                .query(question)
                .similarityThreshold(0.6)
                .topK(3).build();
        List<Document> documents = vectorStore.similaritySearch(searchRequest);
        // 2. 封装提示词
        String template = """
                请基于以下上下⽂回答问题：
                 {documents}
                ⽤户问题：{question}
                """;
        PromptTemplate promptTemplate = new PromptTemplate(template);
        promptTemplate.add("documents", documents.stream()
                .map(Document::getText)
                .collect(Collectors.joining("\n")));
        promptTemplate.add("question", question);
        // 3. 将封装的提示词和问题提交给大模型
        return ollamaChatClient.prompt(promptTemplate.create())
                .user(question)
                .stream()
                .content();
    }
}
